import librosa
import numpy as np


def calculate_audio_psnr(original_audio, watermarked_audio):
    """
    计算音频的PSNR值
    :param original_audio: 原始音频信号
    :param watermarked_audio: 加入水印后的音频信号
    :return: PSNR值
    """
    # 检查两个音频信号的长度是否相同
    if len(original_audio) != len(watermarked_audio):
        raise ValueError("原始音频和加水印后的音频长度必须相同。")

    # 计算均方误差（MSE）
    mse = np.mean((original_audio - watermarked_audio) ** 2)
    if mse == 0:
        return float('inf')

    # 音频信号的最大可能值
    max_audio = 1.0

    # 计算PSNR
    psnr = 20 * np.log10(max_audio / np.sqrt(mse))
    return psnr


# 读取原始音频文件和加入水印后的音频文件
original_file = 'original.wav'
watermarked_file = 'output.wav'

try:
    original_audio, _ = librosa.load(original_file, sr=None)
    watermarked_audio, _ = librosa.load(watermarked_file, sr=None)

    # 计算PSNR
    psnr = calculate_audio_psnr(original_audio, watermarked_audio)
    print(f"音频的PSNR值为: {psnr} dB")

except FileNotFoundError:
    print("文件未找到，请检查文件路径和文件名。")
except ValueError as e:
    print(f"错误: {e}")
